"""
Kafka客户端工具，提供生产者和消费者功能
"""

from confluent_kafka import Producer, Consumer, KafkaError
import json
import logging
import os
import sys
import time
from typing import Dict, List, Any, Callable, Optional

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config.kafka_config import get_producer_config, get_consumer_config, TOPICS

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class KafkaProducer:
    """Kafka生产者类，用于发送消息到Kafka主题"""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None, topic: Optional[str] = None):
        """
        初始化Kafka生产者
        
        参数:
            config (dict, optional): Kafka生产者配置
            topic (str, optional): 默认主题
        """
        self.config = config or get_producer_config()
        self.default_topic = topic
        self.producer = None
        self._connect()
    
    def _connect(self):
        """连接到Kafka服务器"""
        try:
            self.producer = Producer(self.config)
            logger.info(f"Kafka生产者已连接到 {self.config.get('bootstrap_servers', '')}")
        except Exception as e:
            logger.error(f"Kafka生产者连接失败: {e}")
            self.producer = None
    
    def send(self, value: Any, topic: Optional[str] = None, key: Optional[str] = None,
             partition: Optional[int] = None, timestamp: Optional[int] = None, headers: Optional[List[Dict]] = None):
        """
        发送消息到Kafka主题
        
        参数:
            value (Any): 消息内容，如果不是字符串，将被转换为JSON
            topic (str, optional): 目标主题，如果未提供则使用默认主题
            key (str, optional): 消息键
            partition (int, optional): 分区号
            timestamp (int, optional): 消息时间戳
            headers (list, optional): 消息头部
            
        返回:
            bool: 是否成功发送
        """
        if self.producer is None:
            logger.error("Kafka生产者未连接")
            return False
        
        topic = topic or self.default_topic
        if topic is None:
            logger.error("未指定主题")
            return False
        
        # 处理消息值
        try:
            if not isinstance(value, str):
                value = json.dumps(value)
            value = value.encode('utf-8')
        except Exception as e:
            logger.error(f"消息编码失败: {e}")
            return False
        
        # 处理消息键
        if key and not isinstance(key, bytes):
            key = str(key).encode('utf-8')
        
        # 处理消息头部
        kafka_headers = None
        if headers:
            kafka_headers = [(k, v.encode('utf-8') if isinstance(v, str) else v) 
                            for h in headers for k, v in h.items()]
        
        # 发送消息
        try:
            self.producer.produce(
                topic=topic,
                value=value,
                key=key,
                partition=partition,
                timestamp=timestamp,
                headers=kafka_headers
            )
            # 刷新以确保消息被发送
            self.producer.flush(timeout=10)
            return True
        except Exception as e:
            logger.error(f"发送消息失败: {e}")
            return False
    
    def send_batch(self, messages: List[Dict[str, Any]], topic: Optional[str] = None):
        """
        批量发送消息
        
        参数:
            messages (list): 消息字典列表，每个字典可包含'value', 'key', 'partition', 'timestamp', 'headers'
            topic (str, optional): 目标主题，如果未提供则使用默认主题
            
        返回:
            int: 成功发送的消息数量
        """
        if self.producer is None:
            logger.error("Kafka生产者未连接")
            return 0
        
        topic = topic or self.default_topic
        if topic is None:
            logger.error("未指定主题")
            return 0
        
        count = 0
        for msg in messages:
            value = msg.get('value')
            key = msg.get('key')
            partition = msg.get('partition')
            timestamp = msg.get('timestamp')
            headers = msg.get('headers')
            
            if self.send(value, topic, key, partition, timestamp, headers):
                count += 1
        
        return count
    
    def close(self):
        """关闭生产者连接"""
        if self.producer:
            self.producer.flush()
            logger.info("Kafka生产者已关闭")
            self.producer = None

class KafkaConsumer:
    """Kafka消费者类，用于从Kafka主题消费消息"""
    
    def __init__(self, topics: Optional[List[str]] = None, config: Optional[Dict[str, Any]] = None, 
                 group_id: Optional[str] = None, auto_commit: bool = True):
        """
        初始化Kafka消费者
        
        参数:
            topics (list, optional): 要订阅的主题列表
            config (dict, optional): Kafka消费者配置
            group_id (str, optional): 消费者组ID
            auto_commit (bool): 是否自动提交偏移量
        """
        self.config = config or get_consumer_config(group_id)
        if not auto_commit:
            self.config['enable.auto.commit'] = False
        
        self.topics = topics or []
        self.consumer = None
        self.running = False
        self._connect()
    
    def _connect(self):
        """连接到Kafka服务器"""
        try:
            self.consumer = Consumer(self.config)
            if self.topics:
                self.consumer.subscribe(self.topics)
                logger.info(f"Kafka消费者已订阅主题: {', '.join(self.topics)}")
            logger.info(f"Kafka消费者已连接到 {self.config.get('bootstrap.servers', '')}")
        except Exception as e:
            logger.error(f"Kafka消费者连接失败: {e}")
            self.consumer = None
    
    def subscribe(self, topics: List[str]):
        """
        订阅主题
        
        参数:
            topics (list): 要订阅的主题列表
            
        返回:
            bool: 是否成功订阅
        """
        if self.consumer is None:
            logger.error("Kafka消费者未连接")
            return False
        
        try:
            self.consumer.subscribe(topics)
            self.topics = topics
            logger.info(f"已订阅主题: {', '.join(topics)}")
            return True
        except Exception as e:
            logger.error(f"订阅主题失败: {e}")
            return False
    
    def consume_one(self, timeout: float = 1.0):
        """
        消费单条消息
        
        参数:
            timeout (float): 超时时间（秒）
            
        返回:
            dict: 消息内容，None表示超时或错误
        """
        if self.consumer is None:
            logger.error("Kafka消费者未连接")
            return None
        
        try:
            msg = self.consumer.poll(timeout)
            
            if msg is None:
                return None
            
            if msg.error():
                if msg.error().code() == KafkaError._PARTITION_EOF:
                    logger.debug(f"到达分区末尾 {msg.topic()} [{msg.partition()}]")
                else:
                    logger.error(f"消费消息错误: {msg.error()}")
                return None
            
            # 处理消息
            value = msg.value()
            try:
                if value is not None:
                    value = value.decode('utf-8')
                    try:
                        value = json.loads(value)
                    except json.JSONDecodeError:
                        # 非JSON格式，保持字符串
                        pass
            except Exception as e:
                logger.warning(f"解码消息内容失败: {e}")
            
            # 处理消息键
            key = msg.key()
            if key is not None:
                try:
                    key = key.decode('utf-8')
                except Exception:
                    pass
            
            # 处理消息头部
            headers = []
            if msg.headers():
                for header in msg.headers():
                    k, v = header
                    try:
                        v = v.decode('utf-8')
                    except Exception:
                        pass
                    headers.append({k: v})
            
            result = {
                'topic': msg.topic(),
                'partition': msg.partition(),
                'offset': msg.offset(),
                'timestamp': msg.timestamp(),
                'value': value,
                'key': key,
                'headers': headers if headers else None
            }
            
            return result
        except Exception as e:
            logger.error(f"消费消息失败: {e}")
            return None
    
    def consume_batch(self, batch_size: int = 10, timeout: float = 1.0):
        """
        批量消费消息
        
        参数:
            batch_size (int): 批量大小
            timeout (float): 每条消息的超时时间（秒）
            
        返回:
            list: 消息内容列表
        """
        messages = []
        start_time = time.time()
        
        while len(messages) < batch_size and (time.time() - start_time) < (timeout * batch_size):
            msg = self.consume_one(timeout)
            if msg:
                messages.append(msg)
        
        return messages
    
    def consume_loop(self, process_func: Callable, batch_size: int = 1, 
                    timeout: float = 1.0, max_messages: Optional[int] = None,
                    error_handler: Optional[Callable] = None):
        """
        循环消费消息
        
        参数:
            process_func (callable): 处理消息的函数，接收消息作为参数
            batch_size (int): 批量大小，1表示单条处理
            timeout (float): 每条消息的超时时间（秒）
            max_messages (int, optional): 最大消费消息数，None表示无限
            error_handler (callable, optional): 错误处理函数
            
        返回:
            int: 已处理的消息数量
        """
        if self.consumer is None:
            logger.error("Kafka消费者未连接")
            return 0
        
        count = 0
        self.running = True
        
        try:
            while self.running and (max_messages is None or count < max_messages):
                try:
                    if batch_size > 1:
                        messages = self.consume_batch(batch_size, timeout)
                        if messages:
                            process_func(messages)
                            count += len(messages)
                    else:
                        msg = self.consume_one(timeout)
                        if msg:
                            process_func(msg)
                            count += 1
                except Exception as e:
                    logger.error(f"处理消息失败: {e}")
                    if error_handler:
                        error_handler(e)
        except KeyboardInterrupt:
            logger.info("接收到中断信号，停止消费")
        finally:
            self.running = False
        
        return count
    
    def stop(self):
        """停止消费循环"""
        self.running = False
    
    def commit(self):
        """手动提交偏移量"""
        if self.consumer:
            try:
                self.consumer.commit()
                return True
            except Exception as e:
                logger.error(f"提交偏移量失败: {e}")
                return False
        return False
    
    def close(self):
        """关闭消费者连接"""
        if self.consumer:
            self.consumer.close()
            logger.info("Kafka消费者已关闭")
            self.consumer = None
            self.running = False
            return True
        return False 